"""
内存管理模块
用于管理对话历史和上下文信息
"""

import json
import time
from datetime import datetime
from typing import Dict, Any, List, Optional, Union
from dataclasses import dataclass, asdict
import os
from pathlib import Path


@dataclass
class Message:
    """消息类"""
    role: str
    content: str
    timestamp: float
    metadata: Dict[str, Any] = None
    
    def __post_init__(self):
        if self.metadata is None:
            self.metadata = {}
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return asdict(self)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'Message':
        """从字典创建消息"""
        return cls(**data)


class ConversationMemory:
    """对话记忆管理"""
    
    def __init__(self, max_messages: int = 50, max_tokens: int = 4000):
        self.max_messages = max_messages
        self.max_tokens = max_tokens
        self.messages: List[Message] = []
        self.context: Dict[str, Any] = {}
        self.session_id = None
        self.created_at = time.time()
        
    def add_message(self, role: str, content: str, metadata: Dict[str, Any] = None):
        """添加消息到记忆中"""
        message = Message(
            role=role,
            content=content,
            timestamp=time.time(),
            metadata=metadata or {}
        )
        self.messages.append(message)
        
        # 限制消息数量
        if len(self.messages) > self.max_messages:
            self.messages = self.messages[-self.max_messages:]
    
    def get_messages(self, limit: int = None) -> List[Message]:
        """获取消息列表"""
        if limit:
            return self.messages[-limit:]
        return self.messages
    
    def get_formatted_history(self, last_n: int = None) -> str:
        """获取格式化的对话历史"""
        messages = self.get_messages(last_n)
        if not messages:
            return "暂无对话历史"
        
        formatted_lines = []
        for msg in messages:
            role_name = "用户" if msg.role == "user" else "助手"
            formatted_lines.append(f"{role_name}: {msg.content}")
        
        return "\n".join(formatted_lines)
    
    def get_recent_messages(self, minutes: int = 30) -> List[Message]:
        """获取最近的消息"""
        cutoff_time = time.time() - (minutes * 60)
        return [msg for msg in self.messages if msg.timestamp > cutoff_time]
    
    def clear_messages(self):
        """清空消息"""
        self.messages = []
    
    def set_context(self, key: str, value: Any):
        """设置上下文信息"""
        self.context[key] = value
    
    def get_context(self, key: str, default: Any = None) -> Any:
        """获取上下文信息"""
        return self.context.get(key, default)
    
    def clear_context(self):
        """清空上下文"""
        self.context = {}
    
    def get_conversation_summary(self) -> Dict[str, Any]:
        """获取对话摘要"""
        if not self.messages:
            return {
                "total_messages": 0,
                "summary": "暂无对话记录",
                "last_message_time": None,
                "conversation_duration": 0
            }
        
        total_messages = len(self.messages)
        last_message = self.messages[-1]
        first_message = self.messages[0]
        
        # 计算对话持续时间
        duration = last_message.timestamp - first_message.timestamp
        
        # 生成简要摘要
        recent_messages = self.get_messages(5)  # 最近5条消息
        summary_parts = []
        for msg in recent_messages:
            role_name = "用户" if msg.role == "user" else "助手"
            content_preview = msg.content[:50] + "..." if len(msg.content) > 50 else msg.content
            summary_parts.append(f"{role_name}: {content_preview}")
        
        return {
            "total_messages": total_messages,
            "summary": "\n".join(summary_parts),
            "last_message_time": datetime.fromtimestamp(last_message.timestamp).isoformat(),
            "conversation_duration": round(duration, 2)
        }
    
    def clear_history(self):
        """清空对话历史"""
        self.clear_messages()
    
    @property
    def conversation_history(self) -> List[Message]:
        """获取对话历史（属性访问）"""
        return self.messages
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            "session_id": self.session_id,
            "created_at": self.created_at,
            "messages": [msg.to_dict() for msg in self.messages],
            "context": self.context,
            "max_messages": self.max_messages,
            "max_tokens": self.max_tokens
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'ConversationMemory':
        """从字典创建对话记忆"""
        memory = cls(
            max_messages=data.get("max_messages", 50),
            max_tokens=data.get("max_tokens", 4000)
        )
        memory.session_id = data.get("session_id")
        memory.created_at = data.get("created_at", time.time())
        memory.context = data.get("context", {})
        
        # 恢复消息
        for msg_data in data.get("messages", []):
            memory.messages.append(Message.from_dict(msg_data))
        
        return memory


class ContextManager:
    """上下文管理器"""
    
    def __init__(self, storage_path: str = None):
        self.storage_path = storage_path or "conversations"
        self.conversations: Dict[str, ConversationMemory] = {}
        self.current_session_id = None
        
        # 确保存储目录存在
        Path(self.storage_path).mkdir(parents=True, exist_ok=True)
    
    def create_session(self, session_id: str = None) -> str:
        """创建新的会话"""
        if session_id is None:
            session_id = f"session_{int(time.time())}"
        
        if session_id not in self.conversations:
            memory = ConversationMemory()
            memory.session_id = session_id
            self.conversations[session_id] = memory
        
        self.current_session_id = session_id
        return session_id
    
    def get_session(self, session_id: str) -> Optional[ConversationMemory]:
        """获取会话"""
        return self.conversations.get(session_id)
    
    def get_current_session(self) -> Optional[ConversationMemory]:
        """获取当前会话"""
        if self.current_session_id:
            return self.conversations.get(self.current_session_id)
        return None
    
    def switch_session(self, session_id: str) -> bool:
        """切换会话"""
        if session_id in self.conversations:
            self.current_session_id = session_id
            return True
        return False
    
    def delete_session(self, session_id: str) -> bool:
        """删除会话"""
        if session_id in self.conversations:
            del self.conversations[session_id]
            # 删除存储文件
            file_path = Path(self.storage_path) / f"{session_id}.json"
            if file_path.exists():
                file_path.unlink()
            return True
        return False
    
    def list_sessions(self) -> List[str]:
        """列出所有会话"""
        return list(self.conversations.keys())
    
    def save_session(self, session_id: str) -> bool:
        """保存会话到文件"""
        if session_id not in self.conversations:
            return False
        
        try:
            file_path = Path(self.storage_path) / f"{session_id}.json"
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(self.conversations[session_id].to_dict(), f, ensure_ascii=False, indent=2)
            return True
        except Exception as e:
            print(f"保存会话失败: {e}")
            return False
    
    def load_session(self, session_id: str) -> bool:
        """从文件加载会话"""
        try:
            file_path = Path(self.storage_path) / f"{session_id}.json"
            if not file_path.exists():
                return False
            
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            memory = ConversationMemory.from_dict(data)
            self.conversations[session_id] = memory
            return True
        except Exception as e:
            print(f"加载会话失败: {e}")
            return False
    
    def save_all_sessions(self):
        """保存所有会话"""
        for session_id in self.conversations:
            self.save_session(session_id)
    
    def load_all_sessions(self):
        """加载所有会话"""
        storage_path = Path(self.storage_path)
        if not storage_path.exists():
            return
        
        for file_path in storage_path.glob("*.json"):
            session_id = file_path.stem
            self.load_session(session_id)
    
    def cleanup_old_sessions(self, days: int = 30):
        """清理旧会话"""
        cutoff_time = time.time() - (days * 24 * 60 * 60)
        
        sessions_to_delete = []
        for session_id, memory in self.conversations.items():
            if memory.created_at < cutoff_time:
                sessions_to_delete.append(session_id)
        
        for session_id in sessions_to_delete:
            self.delete_session(session_id)
    
    def get_session_stats(self) -> Dict[str, Any]:
        """获取会话统计信息"""
        total_sessions = len(self.conversations)
        total_messages = sum(len(memory.messages) for memory in self.conversations.values())
        
        return {
            "total_sessions": total_sessions,
            "total_messages": total_messages,
            "current_session": self.current_session_id,
            "active_sessions": len([m for m in self.conversations.values() if m.messages])
        }


# 全局实例
conversation_memory = ConversationMemory()
context_manager = ContextManager()

# 便捷函数
def add_message(role: str, content: str, metadata: Dict[str, Any] = None):
    """添加消息到全局记忆"""
    conversation_memory.add_message(role, content, metadata)

def get_messages(limit: int = None) -> List[Message]:
    """获取全局记忆中的消息"""
    return conversation_memory.get_messages(limit)

def clear_memory():
    """清空全局记忆"""
    conversation_memory.clear_messages()
    conversation_memory.clear_context()

def set_context(key: str, value: Any):
    """设置全局上下文"""
    conversation_memory.set_context(key, value)

def get_context(key: str, default: Any = None) -> Any:
    """获取全局上下文"""
    return conversation_memory.get_context(key, default)
